#include "LowLevelCapture.h"
#include "Utils.h"
#include <iostream>
#include <tlhelp32.h>

#pragma comment(lib, "ntdll.lib")

// Forward declarations for NT API functions
extern "C" {
    NTSTATUS NTAPI NtOpenSection(PHANDLE SectionHandle, ACCESS_MASK DesiredAccess, PVOID ObjectAttributes);
    NTSTATUS NTAPI NtMapViewOfSection(HANDLE SectionHandle, HANDLE ProcessHandle, PVOID* BaseAddress,
                                     ULONG_PTR ZeroBits, SIZE_T CommitSize, PVOID SectionOffset,
                                     PSIZE_T ViewSize, DWORD InheritDisposition, ULONG AllocationType, ULONG Win32Protect);
    NTSTATUS NTAPI NtUnmapViewOfSection(HANDLE ProcessHandle, PVOID BaseAddress);
}

LowLevelCapture::LowLevelCapture() :
    m_initialized(false),
    m_hFrameBuffer(NULL),
    m_pFrameBufferData(nullptr) {
}

LowLevelCapture::~LowLevelCapture() {
    Cleanup();
}

bool LowLevelCapture::Initialize() {
    Utils::DebugLog("LowLevelCapture: Initializing low-level capture");

    if (!EnableRequiredPrivileges()) {
        Utils::DebugLog("LowLevelCapture: Warning - Could not enable all required privileges");
    }

    m_initialized = true;
    Utils::DebugLog("LowLevelCapture: Initialized successfully");
    return true;
}

void LowLevelCapture::Cleanup() {
    if (m_pFrameBufferData && m_hFrameBuffer) {
        NtUnmapViewOfSection(GetCurrentProcess(), m_pFrameBufferData);
        m_pFrameBufferData = nullptr;
    }

    if (m_hFrameBuffer) {
        CloseHandle(m_hFrameBuffer);
        m_hFrameBuffer = NULL;
    }

    m_initialized = false;
}

bool LowLevelCapture::EnableRequiredPrivileges() {
    HANDLE hToken;
    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) {
        return false;
    }

    // Enable debug privilege for low-level access
    TOKEN_PRIVILEGES tp;
    tp.PrivilegeCount = 1;
    tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

    bool result = true;

    // Enable SE_DEBUG_NAME
    if (LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &tp.Privileges[0].Luid)) {
        AdjustTokenPrivileges(hToken, FALSE, &tp, 0, NULL, NULL);
    }

    // Enable SE_PROF_SINGLE_PROCESS_NAME if available
    if (LookupPrivilegeValue(NULL, SE_PROF_SINGLE_PROCESS_NAME, &tp.Privileges[0].Luid)) {
        AdjustTokenPrivileges(hToken, FALSE, &tp, 0, NULL, NULL);
    }

    // Enable SE_SYSTEM_PROFILE_NAME if available
    if (LookupPrivilegeValue(NULL, SE_SYSTEM_PROFILE_NAME, &tp.Privileges[0].Luid)) {
        AdjustTokenPrivileges(hToken, FALSE, &tp, 0, NULL, NULL);
    }

    CloseHandle(hToken);
    return result;
}

bool LowLevelCapture::CaptureActiveDesktop(HBITMAP& hBitmap, int& width, int& height) {
    if (!m_initialized) {
        m_lastError = "LowLevelCapture not initialized";
        return false;
    }

    Utils::DebugLog("LowLevelCapture: Attempting low-level desktop capture");

    // Method 1: Try to access frame buffer directly
    if (CaptureWithSystemSession(hBitmap, width, height)) {
        Utils::DebugLog("LowLevelCapture: System session capture succeeded");
        return true;
    }

    // Method 2: Try DirectDraw surface capture
    if (CaptureWithDirectDrawSurface(hBitmap, width, height)) {
        Utils::DebugLog("LowLevelCapture: DirectDraw capture succeeded");
        return true;
    }

    // Method 3: Try mirror driver approach
    if (CaptureWithMirrorDriver(hBitmap, width, height)) {
        Utils::DebugLog("LowLevelCapture: Mirror driver capture succeeded");
        return true;
    }

    m_lastError = "All low-level capture methods failed";
    return false;
}

bool LowLevelCapture::CaptureWithSystemSession(HBITMAP& hBitmap, int& width, int& height) {
    Utils::DebugLog("LowLevelCapture: Trying system session capture");

    // Try to access the console session directly
    DWORD activeConsoleSessionId = WTSGetActiveConsoleSessionId();

    // Get the session's desktop
    HDESK hSessionDesktop = NULL;

    // Try different methods to get session desktop
    HWINSTA hWinStation = OpenWindowStationA("WinSta0", FALSE, GENERIC_READ);
    if (hWinStation) {
        hSessionDesktop = OpenDesktopA("Default", 0, FALSE,
            DESKTOP_READOBJECTS | DESKTOP_WRITEOBJECTS | DESKTOP_CREATEWINDOW);

        if (!hSessionDesktop) {
            hSessionDesktop = OpenDesktopA("Winlogon", 0, FALSE,
                DESKTOP_READOBJECTS | DESKTOP_WRITEOBJECTS);
        }

        CloseWindowStation(hWinStation);
    }

    if (!hSessionDesktop) {
        Utils::DebugLog("LowLevelCapture: Could not open session desktop");
        return false;
    }

    // Switch to session desktop temporarily
    HDESK hOriginalDesktop = GetThreadDesktop(GetCurrentThreadId());
    bool captureSuccess = false;

    if (SetThreadDesktop(hSessionDesktop)) {
        Utils::DebugLog("LowLevelCapture: Switched to session desktop");

        // Now try to capture
        int screenWidth = GetSystemMetrics(SM_CXSCREEN);
        int screenHeight = GetSystemMetrics(SM_CYSCREEN);

        if (screenWidth > 0 && screenHeight > 0) {
            HDC hScreenDC = GetDC(NULL);
            if (hScreenDC) {
                HDC hMemDC = CreateCompatibleDC(hScreenDC);
                if (hMemDC) {
                    HBITMAP hScreenBitmap = CreateCompatibleBitmap(hScreenDC, screenWidth, screenHeight);
                    if (hScreenBitmap) {
                        HGDIOBJ hOldBitmap = SelectObject(hMemDC, hScreenBitmap);

                        // Force a window update first
                        InvalidateRect(NULL, NULL, TRUE);
                        UpdateWindow(GetDesktopWindow());

                        BOOL result = BitBlt(hMemDC, 0, 0, screenWidth, screenHeight,
                                           hScreenDC, 0, 0, SRCCOPY | CAPTUREBLT);

                        SelectObject(hMemDC, hOldBitmap);

                        if (result) {
                            hBitmap = hScreenBitmap;
                            width = screenWidth;
                            height = screenHeight;
                            captureSuccess = true;
                            Utils::DebugLog("LowLevelCapture: Session desktop capture succeeded");
                        } else {
                            DeleteObject(hScreenBitmap);
                            Utils::DebugLog("LowLevelCapture: BitBlt failed: " + Utils::GetLastWin32Error());
                        }
                    }
                    DeleteDC(hMemDC);
                }
                ReleaseDC(NULL, hScreenDC);
            }
        }

        // Restore original desktop
        SetThreadDesktop(hOriginalDesktop);
    }

    CloseDesktop(hSessionDesktop);
    return captureSuccess;
}

bool LowLevelCapture::CaptureWithDirectDrawSurface(HBITMAP& hBitmap, int& width, int& height) {
    Utils::DebugLog("LowLevelCapture: Trying DirectDraw surface capture");

    // This method uses a different approach - accessing the primary surface
    int screenWidth = GetSystemMetrics(SM_CXSCREEN);
    int screenHeight = GetSystemMetrics(SM_CYSCREEN);

    if (screenWidth <= 0 || screenHeight <= 0) {
        return false;
    }

    // Create device context for the entire screen
    HDC hScreenDC = CreateDCA("DISPLAY", NULL, NULL, NULL);
    if (!hScreenDC) {
        return false;
    }

    HDC hMemDC = CreateCompatibleDC(hScreenDC);
    if (!hMemDC) {
        DeleteDC(hScreenDC);
        return false;
    }

    HBITMAP hScreenBitmap = CreateCompatibleBitmap(hScreenDC, screenWidth, screenHeight);
    if (!hScreenBitmap) {
        DeleteDC(hMemDC);
        DeleteDC(hScreenDC);
        return false;
    }

    HGDIOBJ hOldBitmap = SelectObject(hMemDC, hScreenBitmap);

    // Try multiple capture methods
    BOOL result = FALSE;

    // Method 1: Standard BitBlt
    result = BitBlt(hMemDC, 0, 0, screenWidth, screenHeight, hScreenDC, 0, 0, SRCCOPY);

    if (!result) {
        // Method 2: BitBlt with CAPTUREBLT
        result = BitBlt(hMemDC, 0, 0, screenWidth, screenHeight, hScreenDC, 0, 0, SRCCOPY | CAPTUREBLT);
    }

    if (!result) {
        // Method 3: StretchBlt
        result = StretchBlt(hMemDC, 0, 0, screenWidth, screenHeight,
                           hScreenDC, 0, 0, screenWidth, screenHeight, SRCCOPY);
    }

    SelectObject(hMemDC, hOldBitmap);
    DeleteDC(hMemDC);
    DeleteDC(hScreenDC);

    if (result) {
        hBitmap = hScreenBitmap;
        width = screenWidth;
        height = screenHeight;
        Utils::DebugLog("LowLevelCapture: DirectDraw surface capture succeeded");
        return true;
    } else {
        DeleteObject(hScreenBitmap);
        return false;
    }
}

bool LowLevelCapture::CaptureWithMirrorDriver(HBITMAP& hBitmap, int& width, int& height) {
    Utils::DebugLog("LowLevelCapture: Trying mirror driver approach");

    // This is a placeholder for mirror driver implementation
    // Mirror drivers require kernel-mode components and are complex to implement
    // For now, we'll use a simulated approach that tries to access video memory directly

    return false; // Not implemented in this version
}